package com.example.config.loadbalanced;

import com.google.common.collect.Sets;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.HttpRequest;
import org.springframework.http.client.ClientHttpRequestExecution;
import org.springframework.http.client.ClientHttpRequestInterceptor;
import org.springframework.http.client.ClientHttpResponse;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.context.SecurityContextHolder;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.Cookie;
import javax.servlet.http.HttpServletRequest;
import java.io.IOException;
import java.util.Arrays;
import java.util.Enumeration;
import java.util.Set;

@Component
public class EnhancedTokenInterceptor implements ClientHttpRequestInterceptor {

    @Autowired(required = false)
    private TokenProvider tokenProvider; // 自定义的 token 提供器

    @Override
    public ClientHttpResponse intercept(HttpRequest request, byte[] body, ClientHttpRequestExecution execution)
            throws IOException {

        // 多种方式获取 token（按优先级）
        String token = getTokenWithPriority();

        if (StringUtils.hasText(token)) {
            request.getHeaders().set("Authorization", "Bearer " + token);
        }

        // 复制所有相关的头信息
        copyHeadersFromCurrentRequest(request);

        return execution.execute(request, body);
    }

    private String getTokenWithPriority() {
        // 1. 从自定义 TokenProvider 获取
        if (tokenProvider != null) {
            String token = tokenProvider.getCurrentToken();
            if (StringUtils.hasText(token)) {
                return token;
            }
        }

        // 2. 从 Spring Security Context 获取
        Authentication authentication = SecurityContextHolder.getContext().getAuthentication();
        if (authentication != null) {
            Object credentials = authentication.getCredentials();
            if (credentials instanceof String && StringUtils.hasText((String) credentials)) {
                return (String) credentials;
            }
        }
        // 3. 从当前 HTTP 请求获取
        ServletRequestAttributes attributes = (ServletRequestAttributes)
                RequestContextHolder.getRequestAttributes();
        if (attributes != null) {
            HttpServletRequest currentRequest = attributes.getRequest();
            return extractTokenFromRequest(currentRequest);
        }

        return null;
    }

    private String extractTokenFromRequest(HttpServletRequest request) {
        // 从 Authorization 头获取
        String authHeader = request.getHeader("Authorization");
        if (StringUtils.hasText(authHeader) && authHeader.startsWith("Bearer ")) {
            return authHeader.substring(7);
        }

        // 从 query parameter 获取
        String tokenParam = request.getParameter("token");
        if (StringUtils.hasText(tokenParam)) {
            return tokenParam;
        }

        // 从 cookie 获取
        Cookie[] cookies = request.getCookies();
        if (cookies != null) {
            for (Cookie cookie : cookies) {
                if ("access_token".equals(cookie.getName())) {
                    return cookie.getValue();
                }
            }
        }

        return null;
    }

    private void copyHeadersFromCurrentRequest(HttpRequest targetRequest) {
        ServletRequestAttributes attributes = (ServletRequestAttributes)
                RequestContextHolder.getRequestAttributes();
        if (attributes == null) return;

        HttpServletRequest currentRequest = attributes.getRequest();
        Enumeration<String> headerNames = currentRequest.getHeaderNames();

        // 需要转发的头信息
        Set<String> forwardHeaders = Sets.newConcurrentHashSet(Arrays.asList(
                "X-User-Id", "X-User-Name", "X-Roles", "X-Tenant-Id",
                "X-Request-ID", "X-Forwarded-For", "User-Agent")
        );

        while (headerNames.hasMoreElements()) {
            String headerName = headerNames.nextElement();
            if (forwardHeaders.contains(headerName)) {
                String headerValue = currentRequest.getHeader(headerName);
                targetRequest.getHeaders().add(headerName, headerValue);
            }
        }
    }
}
